# app.py
from flask import Flask, request, jsonify
import detect  # 引入 detect.py 文件
import io

app = Flask(__name__)

# 加载YOLOv5模型路径
MODEL_PATH = "D:\\yolo\\yolov5-master\\yolov5-master\\runs\\train\\exp5\\weights/best.pt"
model = detect.load_model(MODEL_PATH)  # 加载自定义模型

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        # 确保请求中包含图像文件
        if 'image' not in request.files:
            return jsonify({'error': 'No image part'}), 400
        
        image_file = request.files['image']
        
        if image_file.filename == '':
            return jsonify({'error': 'No selected file'}), 400
        
        # 读取图像文件
        image_bytes = image_file.read()
        
        # 使用 detect.py 中的检测函数
        prediction = detect.detect_image(model, image_bytes)
        
        # 返回 JSON 格式的检测结果
        return jsonify(prediction)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)
